############
# Balance tables
############
rm(list=ls())
library(foreign)
library(survey)

# The dataset is based on the output from the ensemble estimation,
# but we had to convert the CSV to DTA because the file encoding was
# unreadable (due to spanish characters).  

coldat <- read.dta("COLOMBIA_STEP9_PSCORE_IMP1.dta")

# Covariates

controls_reint_ix <- c("index_reint_prog_dissatis",
					"index_reint_state",
					"index_reint_prog_exposure")

controls_conf_ix <- c("index_con_time",
					"index_con_rank_mid_commander",
					"index_con_rank_top_commander",
					"index_con_role_combatant",
					"index_con_unit_cohesion",
					"index_con_unit_discipline",
					"index_con_unit_hierarchy",
					"index_con_unit_motivation") 

controls_demob_ix <- c("index_demob_coerced",
					"index_demob_type") 

controls_join_ix <- c("index_join_ed",
					"index_join_fam",
					"index_join_forced",
					"index_join_griev",
					"index_join_ideo",
					"index_join_material",
					"index_join_network",
					"index_join_psych",
					"index_join_security",
					"index_join_welfare")

controls_gen <- c(  "p8_sex",
					"p11_race_REC1",
					"p9_age",
					"p104_minor",
					"p143_disabled_REC1",
					"p25_grp_demob_name_AUC",
					"p122_risk_all",
					"p121_time_all",
					"p19_ageenter_group1")

X.names <-c(controls_reint_ix,
			controls_conf_ix,
			controls_demob_ix,
			controls_join_ix,
			controls_gen)

X.names.fmt <- c("Index: reint. prog. dissatis.", 
				"Index: reint. formal",
				"Index: reint. prog. exposure",
				"Index: conflict time",
				"Index: mid rank experience",
				"Index: top rank experience",
				"Index: role combatant",
				"Index: unit cohesion",
				"Index: unit discipline",
				"Index: unit hierarchy",
				"Index: unit motivation",
				"Index: coerced demob.",
				"Index: type of demob.",
				"Index: pre-join ed.",
				"Index: pre-join family",
				"Index: forced recruit",
				"Index: join for grievance",
				"Index: join for ideol.",
				"Index: join for material",
				"Index: join for network",
				"Index: join for psych.",
				"Index: join for security",
				"Index: join for welfare",
				"Male",
				"Non-white",
				"Age",
				"Minor when joined",
				"Disabled",
				"AUC",
				"Risk aversion",
				"Time in group",
				"Age when joined")

int.vars.fmt <- c("Employment",
				"Security",
				"Confidence",
				"Depression",
				"Excom. peers",
				"Ties to commander")

# We want all of these to be standardized variables
stanD <- function(x) (x - mean(x))/sd(x)
for(j in 1:length(X.names)){
coldat[,X.names[j]] <- stanD(coldat[,X.names[j]])
}

# Fix weight for the one non-participant respondent
coldat$wgt_main <- as.numeric(coldat$wgt_main)
coldat$wgt_main[is.na(coldat$wgt_main)] <- 1

# Placebo tests
int.vars <- names(coldat)[substr(names(coldat), 1, 4)=="int_"]
fitList <- matrix(NA, nrow=length(X.names), ncol=2*length(int.vars))
colnames(fitList) <- c(paste(c("b"), int.vars, sep="_"),paste(c("se"), int.vars, sep="_"))
rownames(fitList) <- X.names

for(j in 1:length(X.names)){

for(i in 1:length(int.vars)){
	#i <- 1

	# Construct D
	Aup <- int.vars[i]
	# To deal with the truncated variable names

	g.hat 		<- coldat[,paste("p_",Aup, sep="")]      
	I.A 		<- coldat[,Aup]       
	Yup <- X.names[j]
	Y <- coldat[,Yup]             

	D.ipw <- ((I.A/g.hat)-1)*Y
	coldat[,paste("Dipw_",Aup, sep="")] <- D.ipw
}


colsvy <- svydesign(	ids = ~svypsu,
						strata = ~svystrat,
						weights=~wgt_main,
						data=coldat,
						nest=T)

for(i in 1:length(int.vars)){
	# Compute effect estimates
	Aup <- int.vars[i]
	formipw <- as.formula(paste("~Dipw_",Aup, sep=""))
	psi.ipw <- svymean(formipw, colsvy, na.rm=T)
	fitList[j, paste("b",int.vars[i], sep="_")] <- psi.ipw[1]
	fitList[j, paste("se",int.vars[i], sep="_")] <- sqrt(attributes(psi.ipw)$var)
}
}

nVars <- nrow(fitList)
pdf("ml-balance-post.pdf", 
		width=10, height=10)
par(mar=(c(0,0,0,0)))
par(mfrow=c(1, 1))
plot(c(-.55, (1/4)+length(int.vars)/2), 
		c(-1, nVars+1), 
		type="n")
#abline(v=0)
text(rep(0, nVars), 1:nVars,
		X.names.fmt, #rownames(fitList),
		pos=2, cex=.75)
for(i in 1:length(int.vars)){
#i <- 1
posUp <- i/2
fitUp <- fitList[,paste(c("b","se"),int.vars[i], sep="_")]
points(posUp+fitUp[,1], 1:nVars, pch=19, cex=.5)
segments(	posUp+fitUp[,1]-1.96*fitUp[,2],
			1:nVars,
			posUp+fitUp[,1]+1.96*fitUp[,2],
			1:nVars)
segments(	posUp+fitUp[,1]-1.64*fitUp[,2],
			1:nVars,
			posUp+fitUp[,1]+1.64*fitUp[,2],
			1:nVars,
			lwd=2)
segments(posUp, 0, posUp, nVars)
segments(posUp-.2, 0, posUp+.2,0)
text(c(posUp-.2, posUp-.1, posUp, posUp+.1, posUp+.2), 
		c(-.5,-.5,-.5,-.5,-.5),
		c(-0.2,-0.1,0,0.1,0.2),
		cex=.75)
text(posUp, nVars+1, int.vars.fmt[i], cex=.75)		
}
dev.off()


# Unweighted tests
int.vars <- names(coldat)[substr(names(coldat), 1, 4)=="int_"]
fitList <- matrix(NA, nrow=length(X.names), ncol=2*length(int.vars))
colnames(fitList) <- c(paste(c("b"), int.vars, sep="_"),paste(c("se"), int.vars, sep="_"))
rownames(fitList) <- X.names

for(j in 1:length(X.names)){

for(i in 1:length(int.vars)){
	#i <- 1

	# Construct D
	Aup <- int.vars[i]
	# To deal with the truncated variable names
	g.hat 		<- rep(mean(I.A), nrow(coldat))      
	I.A 		<- coldat[,Aup]       
	Yup <- X.names[j]
	Y <- coldat[,Yup]             

	D.ipw <- ((I.A/g.hat)-1)*Y
	coldat[,paste("Dipw_",Aup, sep="")] <- D.ipw
}


colsvy <- svydesign(	ids = ~svypsu,
						strata = ~svystrat,
						weights=~wgt_main,
						data=coldat,
						nest=T)

for(i in 1:length(int.vars)){
	# Compute effect estimates
	Aup <- int.vars[i]
	formipw <- as.formula(paste("~Dipw_",Aup, sep=""))
	psi.ipw <- svymean(formipw, colsvy, na.rm=T)
	fitList[j, paste("b",int.vars[i], sep="_")] <- psi.ipw[1]
	fitList[j, paste("se",int.vars[i], sep="_")] <- sqrt(attributes(psi.ipw)$var)
}
}

nVars <- nrow(fitList)

pdf("ml-balance-pre.pdf", 
		width=10, height=10)
par(mar=(c(0,0,0,0)))
par(mfrow=c(1, 1))
plot(c(-.55, (1/4)+length(int.vars)/2), 
		c(-1, nVars+1), 
		type="n")
#abline(v=0)
text(rep(0, nVars), 1:nVars,
		X.names.fmt,#rownames(fitList),
		pos=2, cex=.75)
for(i in 1:length(int.vars)){
#i <- 1
posUp <- i/2
fitUp <- fitList[,paste(c("b","se"),int.vars[i], sep="_")]
points(posUp+fitUp[,1], 1:nVars, pch=19, cex=.5)
segments(	posUp+fitUp[,1]-1.96*fitUp[,2],
			1:nVars,
			posUp+fitUp[,1]+1.96*fitUp[,2],
			1:nVars)
segments(	posUp+fitUp[,1]-1.64*fitUp[,2],
			1:nVars,
			posUp+fitUp[,1]+1.64*fitUp[,2],
			1:nVars,
			lwd=2)
segments(posUp, 0, posUp, nVars)
segments(posUp-.2, 0, posUp+.2,0)
text(c(posUp-.2, posUp-.1, posUp, posUp+.1, posUp+.2), 
		c(-.5,-.5,-.5,-.5,-.5),
		c(-0.2,-0.1,0,0.1,0.2),
		cex=.65)
text(posUp, nVars+1, int.vars.fmt[i], cex=.75)		
}
dev.off()
